﻿
import matplotlib.pyplot as plt
from numpy import array, zeros, arange, pi
from scipy import random
from math import sqrt,  exp
import copy
import psyco

psyco.full()

class Model:
    """ base class for model """
    def __init__(self, init, n):
        self.initial_condtions = init
        # problem order
        self.n = n

class SampleModel(Model):
    """ test """
    def F(self, t, y):
        F = zeros((2))
        F[0] = y[1]
        F[1] = -9.80665 + 65.351e-3 * y[1]**2 * exp(-10.53e-5*y[0])

        return F

class Arenstorf(Model):
    def F(self, t, y):
        F = zeros(4)

        mu = 0.012277471
        mu_ = 1. - mu
        D1 = ( (y[2] + mu)**2 + y[3]**2 )**(3./2)
        D2 = ( (y[2] - mu_)**2 + y[3]**2 )**(3./2)

        F[0] = y[2] + 2.*y[1] - mu_*(y[2]+mu)/D1 - mu*(y[2]-mu_)/D2
        F[1] = y[3] - 2.*y[0] - mu_*y[3]/D1 - mu*y[3]/D2
        F[2] = y[0]
        F[3] = y[1]

        return F

class Brusellator(Model):
    def F(self, t, y):
        F = zeros((2))
        F[0] = 2. + (y[0]**2)*y[1] - 9.533*y[0]
        F[1] = 8.533*y[0] - (y[0]**2)*y[1]

        return F

b0 = 8.4967176264065358e-012
b2 = 1.4999999999999999e-013
a0 = 10080.213382055961
a2 = 350.94867139442778
a4 = pi

K = sqrt(b0/a0)
Tau = sqrt(b2/b0)
T = sqrt(sqrt(a4/a0))
Etta = sqrt(a2/(4*sqrt(a0*a4))+ 0.5)

class StochasticProcess(Model):
    def F(self, t, y):
        x = random.normal(0.,1/0.01)
        print x

        dX = (K*x - y[0])/T
        dA = dX - 2*Etta*y[1] + (1/T)*(x-y[2])
        dY = y[1]

        return array([dX, dA, dY])


class Integrator:
    """ base class for integrator """
    def __init__(self, model, tol, h, t_start, t_stop, dense_points):
        self.model = model
        self.tol = tol
        self.h = h
        self.t_stop = t_stop
        self.t_start = t_start
        self.T = []
        self.Y = []
        self.position_in_dense = 1
        self.stop = False
        self.dense_points = dense_points
        if( len(self.dense_points) == 0): raise Exception("Your dense points array has 0 length.")

    def densify(self, t, y0, y1, dy0, dy1):
        self.T.append(t)
        interp_value = (1.-t)*y0 + t*y1 + t*(t-1.)*( (1-2.*t)*(y1-y0) + (t-1.)*self.h*dy0 + t*self.h*dy1 )
        self.Y.append(interp_value)

    def correct_h(self, e, t):

        if e != 0.:
            h_next = 0.9*self.h*(self.tol/e)**0.2
        else:
            h_next = self.h

        # last step ?
        if (self.h > 0.) and ((t+h_next)>=self.t_stop):
            next_h = self.t_stop - t
            self.stop = True
        self.h = h_next

    def rms_error(self, delta):
        return sqrt(sum(delta*delta)/self.model.n)

        self.h = h_next

# TODO: CLEAN!!!
class Adams4(Integrator):
    def __init__(self, model, tol, h, t_start, t_stop, dense_points):
        Integrator.__init__(self, model, tol, h, t_start, t_stop, dense_points)
        if( len(self.dense_points) < 3 ): raise Exception("Dense points array must be at least 3 items long for Adams–Bashforth–Moulton method.")

    def integrate(self):
        # first we have to get y[1] y[2] y[3] using Runge-Kutta 4(5)
        rk4 = RungeKutta5(self.model, self.tol, self.h, self.t_start, self.t_stop, self.dense_points[0:4])

        # inner data
        aT, aY = rk4.integrate()
        # outer data
        self.T = copy.copy(aT)
        self.Y = copy.copy(aY)

        prev_corr_y = copy.copy(self.Y[-1])

        self.t = copy.copy(self.T[-1])
        self.y = copy.copy(self.Y[-1])

        i = 3
        self.position_in_dense = 4
        prev_y = 0
        prev_dy = 0

        while( self.t + self.h < self.t_stop):
            # predictor
            pred_y = aY[i] + (self.h/24.)*(-9.*self.model.F(aT[i-3], aY[i-3]) \
            + 37.*  self.model.F(aT[i-2], aY[i-2]) \
            - 59.*  self.model.F(aT[i-1], aY[i-1]) \
            + 55.*  self.model.F(aT[i], aY[i]))

            # corrector
            corr_y = aY[i] + (self.h/24.)*(self.model.F(aT[i-2], aY[i-2]) \
            - 5.*   self.model.F(aT[i-1], aY[i-1]) \
            + 19.*  self.model.F(aT[i], aY[i])
            + 9.*   self.model.F(self.t+self.h, pred_y))

            # error correction and dense output
            e = self.rms_error(pred_y-corr_y)
            dy = self.model.F(self.t, corr_y)

            self.t += self.h

            aY.append(corr_y)
            aT.append(self.t)

            print 't:' + self.t.__str__() + '\ty:' + corr_y.__str__() + "\t" + self.h.__str__()

            i += 1

            prev_y = copy.copy(corr_y)
            prev_dy = copy.copy(dy)

            if( self.stop ): break

        return aT, aY

class RungeKutta5 (Integrator):
    """ adaptive Runge-Kutta 4(5) integrator """

    def __init__(self, model, tol, h, t_start, t_stop, dense_points):
        Integrator.__init__(self, model, tol, h, t_start, t_stop, dense_points)
        self.C = array([5179./57600., 0., 7571./16695., 393./640.,-92097./339200., 187/2100, 1.0/40.])
        self.D = array([35./384., 0., 500./1113., 125./192.,-2187./6784., 11./84., 0])
        self.K = zeros((7, self.model.n))


    def integrate(self):
        t = self.t_start
        y = self.model.initial_condtions

        # initial values
        self.T.append(t)
        self.Y.append(y)

        prev_t = 0
        prev_y = 0
        prev_dy = 0

        while(not self.stop):
            self.K[0] = self.h*self.model.F(t, y)

            self.K[1] = self.h*self.model.F(t + 1./5 *self.h, y + 1./5*        self.K[0])

            self.K[2] = self.h*self.model.F(t + 3./10*self.h, y + 3./40*       self.K[0] +
             9./40*           self.K[1])

            self.K[3] = self.h*self.model.F(t + 4./5* self.h, y +
            44./45*      self.K[0]
            + -56./15*         self.K[1]
            + 32./9*       self.K[2])

            self.K[4] = self.h*self.model.F(t + 8./9* self.h, y +
            19372./6561* self.K[0]
            + -25360./2187*    self.K[1] +
            64448./6561* self.K[2]
            + -212./729*   self.K[3])

            self.K[5] = self.h*self.model.F(t + self.h, y +
            9017./3168*        self.K[0] +
            -355./33*        self.K[1] +
            46732./5247* self.K[2] +
            49./176*     self.K[3] +
            -5103./18656*self.K[4])

            self.K[6] = self.h*self.model.F(t + self.h, y
            + 35./384*           self.K[0]
            + 0*               self.K[1]
            + 500./1113*   self.K[2]
            + 125./192*    self.K[3]
            + -2187./6784* self.K[4]
            + 11/84*self.K[5])

            E = zeros((self.model.n))
            dy = zeros((self.model.n))

            for i in range(7):
                dy = dy + self.C[i]*self.K[i]
                E = E +(self.C[i]- self.D[i])*self.K[i]

            # RMS error (rounded mean square not Richard Stallman ;)
            e = self.rms_error(E)

            if( e <= self.tol ):
                # if interpolation poit is btwn t and t+1
                if( t <= self.dense_points[self.position_in_dense] < t + self.h ):
                    self.densify(self.dense_points[self.position_in_dense], prev_y, y, prev_dy, dy )
                    self.position_in_dense += 1
                    if(self.position_in_dense == len(self.dense_points) ):
                        self.stop = True
                else:
                    prev_t = copy.copy(t)
                    prev_y = copy.copy(y)
                    prev_dy = copy.copy(dy)

                    t += self.h
                    y += dy

            self.correct_h(e,t)

            if self.stop:
                break


        # convert list to array and return dett skit
        return self.T, self.Y

''' helpers '''

def get_indexed_by(data,i, every_n=100):
    matr = []
    counter = 1
    for k in data:
        if( counter == 100):
            matr.append(float(k[i]))
            counter == 1
        counter += 1
    return matr

def printSoln(X,Y,freq):

    def printHead(n):
        print "\n        x  ",
        for i in range (n):
            print "      y[",i,"] ",
        print

    def printLine(x,y,n):
        print "%13.4e"% x,
        for i in range (n):
            print "%13.4e"% y[i],
        print

    m = len(Y)
    try: n = len(Y[0])
    except TypeError: n = 1
    if freq == 0: freq = m
    printHead(n)
    for i in range(0,m,freq):
        printLine(X[i],Y[i],n)
    if i != m - 1: printLine(X[m - 1],Y[m - 1],n)


TimeStart = 0
TimeEnd = 100
Step = 0.01

Timespace = arange(TimeStart, TimeEnd, Step)

process = StochasticProcess([0,0,0], 3)
integrator = RungeKutta5(process, 0.01, 0.1, TimeStart, TimeEnd, Timespace)

integrator.integrate()

plt.plot(integrator.T,integrator .Y[:, 2])
plt.show()
